#include <mcp/transport/transport_security.h>
#include <mcp/error.h>
#include <mcp/json/value.h>
#include <iostream>
#include <sstream>
#include <chrono>
#include <random>

namespace mcp {
namespace transport {

// SecurityManager implementation
struct SecurityManager::Impl {
    SecurityConfig config;
    std::unordered_map<std::string, std::chrono::steady_clock::time_point> blocked_ips;
    std::unordered_map<std::string, RateLimitInfo> rate_limits;
    mutable std::mutex security_mutex;
    std::chrono::steady_clock::time_point last_cleanup;
    std::mt19937 rng;

    Impl(const SecurityConfig& cfg)
        : config(cfg), last_cleanup(std::chrono::steady_clock::now()),
          rng(std::random_device{}()) {}
};

SecurityManager::SecurityManager(const SecurityConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

SecurityManager::~SecurityManager() = default;

void SecurityManager::set_config(const SecurityConfig& config) {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);
    impl_->config = config;
}

SecurityConfig SecurityManager::get_config() const {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);
    return impl_->config;
}

SecurityResult SecurityManager::validate_connection(const std::string& client_address,
                                                   const std::optional<std::string>& auth_token) {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);

    // Clean up old entries periodically
    cleanup_expired_entries();

    // Check if IP is blocked
    if (is_ip_blocked(client_address)) {
        return SecurityResult{
            .allowed = false,
            .reason = "IP address is blocked",
            .severity = SecuritySeverity::Block
        };
    }

    // Check rate limiting
    auto rate_limit_result = check_rate_limit(client_address);
    if (!rate_limit_result.allowed) {
        return rate_limit_result;
    }

    // Check authentication if required
    if (impl_->config.require_authentication) {
        if (!auth_token) {
            return SecurityResult{
                .allowed = false,
                .reason = "Authentication token required",
                .severity = SecuritySeverity::Authenticate
            };
        }

        auto auth_result = validate_auth_token(*auth_token);
        if (!auth_result.allowed) {
            return auth_result;
        }
    }

    // Check TLS requirement
    if (impl_->config.require_tls && !is_connection_secure()) {
        return SecurityResult{
            .allowed = false,
            .reason = "TLS encryption required",
            .severity = SecuritySeverity::Encrypt
        };
    }

    // Check allowed origins
    if (!impl_->config.allowed_origins.empty()) {
        if (!is_origin_allowed(client_address)) {
            return SecurityResult{
                .allowed = false,
                .reason = "Origin not allowed",
                .severity = SecuritySeverity::Block
            };
        }
    }

    return SecurityResult{
        .allowed = true,
        .reason = "Connection validated",
        .severity = SecuritySeverity::Allow
    };
}

SecurityResult SecurityManager::validate_message(const json::Value& message,
                                               const std::string& client_address) {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);

    // Check message size
    auto size_result = check_message_size(message);
    if (!size_result.allowed) {
        return size_result;
    }

    // Check rate limiting for messages
    auto message_rate_result = check_message_rate_limit(client_address);
    if (!message_rate_result.allowed) {
        return message_rate_result;
    }

    // Check for malicious patterns
    auto pattern_result = check_malicious_patterns(message);
    if (!pattern_result.allowed) {
        return pattern_result;
    }

    return SecurityResult{
        .allowed = true,
        .reason = "Message validated",
        .severity = SecuritySeverity::Allow
    };
}

void SecurityManager::block_ip(const std::string& ip_address,
                              const std::optional<std::chrono::seconds>& duration) {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);

    auto block_until = std::chrono::steady_clock::now();
    if (duration) {
        block_until += *duration;
    } else {
        // Default block duration from config
        block_until += std::chrono::seconds(impl_->config.ip_block_duration);
    }

    impl_->blocked_ips[ip_address] = block_until;

    std::cout << "[Security] Blocked IP: " << ip_address
              << " until " << std::chrono::system_clock::time_point(std::chrono::duration_cast<std::chrono::system_clock::duration>(block_until.time_since_epoch()))
              << std::endl;
}

void SecurityManager::unblock_ip(const std::string& ip_address) {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);
    impl_->blocked_ips.erase(ip_address);
}

bool SecurityManager::is_ip_blocked(const std::string& ip_address) const {
    auto it = impl_->blocked_ips.find(ip_address);
    if (it == impl_->blocked_ips.end()) {
        return false;
    }

    // Check if block has expired
    if (std::chrono::steady_clock::now() > it->second) {
        impl_->blocked_ips.erase(it);
        return false;
    }

    return true;
}

std::vector<std::string> SecurityManager::get_blocked_ips() const {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);

    std::vector<std::string> blocked;
    auto now = std::chrono::steady_clock::now();

    for (auto it = impl_->blocked_ips.begin(); it != impl_->blocked_ips.end();) {
        if (now > it->second) {
            it = impl_->blocked_ips.erase(it);
        } else {
            blocked.push_back(it->first);
            ++it;
        }
    }

    return blocked;
}

SecurityResult SecurityManager::validate_auth_token(const std::string& token) {
    if (token.empty()) {
        return SecurityResult{
            .allowed = false,
            .reason = "Empty authentication token",
            .severity = SecuritySeverity::Authenticate
        };
    }

    // Check against valid tokens
    if (!impl_->config.valid_tokens.empty()) {
        if (impl_->config.valid_tokens.find(token) == impl_->config.valid_tokens.end()) {
            return SecurityResult{
                .allowed = false,
                .reason = "Invalid authentication token",
                .severity = SecuritySeverity::Authenticate
            };
        }
    }

    // Check token format (basic validation)
    if (impl_->config.token_validation) {
        auto format_result = check_token_format(token);
        if (!format_result.allowed) {
            return format_result;
        }
    }

    return SecurityResult{
        .allowed = true,
        .reason = "Authentication token valid",
        .severity = SecuritySeverity::Allow
    };
}

SecurityResult SecurityManager::check_rate_limit(const std::string& client_address) {
    auto now = std::chrono::steady_clock::now();
    auto& info = impl_->rate_limits[client_address];

    // Reset counter if window has passed
    if (now - info.window_start > std::chrono::seconds(impl_->config.rate_limit_window)) {
        info.request_count = 0;
        info.window_start = now;
    }

    // Check if limit exceeded
    if (info.request_count >= impl_->config.rate_limit_max_requests) {
        // Apply escalating penalties
        if (info.violation_count >= impl_->config.rate_limit_violations_before_block) {
            block_ip(client_address);
            return SecurityResult{
                .allowed = false,
                .reason = "Rate limit violations exceeded - IP blocked",
                .severity = SecuritySeverity::Block
            };
        }

        info.violation_count++;
        return SecurityResult{
            .allowed = false,
            .reason = "Rate limit exceeded",
            .severity = SecuritySeverity::Throttle
        };
    }

    info.request_count++;
    return SecurityResult{
        .allowed = true,
        .reason = "Rate limit check passed",
        .severity = SecuritySeverity::Allow
    };
}

SecurityResult SecurityManager::check_message_rate_limit(const std::string& client_address) {
    // Similar to check_rate_limit but with potentially different limits for messages
    return check_rate_limit(client_address); // For simplicity, use same limits
}

SecurityResult SecurityManager::check_message_size(const json::Value& message) {
    // Calculate message size in bytes
    std::string serialized = message.serialize();
    size_t size_bytes = serialized.size();

    if (size_bytes > impl_->config.max_message_size) {
        // Check for size violation escalation
        if (size_bytes > impl_->config.max_message_size * 2) {
            return SecurityResult{
                .allowed = false,
                .reason = "Message size severely exceeds limit",
                .severity = SecuritySeverity::Block
            };
        }

        return SecurityResult{
            .allowed = false,
            .reason = "Message size exceeds limit",
            .severity = SecuritySeverity::Throttle
        };
    }

    return SecurityResult{
        .allowed = true,
        .reason = "Message size acceptable",
        .severity = SecuritySeverity::Allow
    };
}

SecurityResult SecurityManager::check_malicious_patterns(const json::Value& message) {
    std::string serialized = message.serialize();

    // Check for common attack patterns
    for (const auto& pattern : impl_->config.blocked_patterns) {
        if (serialized.find(pattern) != std::string::npos) {
            return SecurityResult{
                .allowed = false,
                .reason = "Message contains blocked pattern",
                .severity = SecuritySeverity::Block
            };
        }
    }

    // Check for JSON injection attempts
    if (serialized.find("<script>") != std::string::npos ||
        serialized.find("javascript:") != std::string::npos) {
        return SecurityResult{
            .allowed = false,
            .reason = "Potential script injection detected",
            .severity = SecuritySeverity::Block
        };
    }

    return SecurityResult{
        .allowed = true,
        .reason = "No malicious patterns detected",
        .severity = SecuritySeverity::Allow
    };
}

SecurityResult SecurityManager::check_token_format(const std::string& token) {
    // Basic token format validation
    if (token.size() < impl_->config.min_token_length) {
        return SecurityResult{
            .allowed = false,
            .reason = "Token too short",
            .severity = SecuritySeverity::Authenticate
        };
    }

    if (token.size() > impl_->config.max_token_length) {
        return SecurityResult{
            .allowed = false,
            .reason = "Token too long",
            .severity = SecuritySeverity::Authenticate
        };
    }

    // Check for valid characters (alphanumeric, hyphen, underscore, period)
    for (char c : token) {
        if (!std::isalnum(c) && c != '-' && c != '_' && c != '.') {
            return SecurityResult{
                .allowed = false,
                .reason = "Token contains invalid characters",
                .severity = SecuritySeverity::Authenticate
            };
        }
    }

    return SecurityResult{
        .allowed = true,
        .reason = "Token format valid",
        .severity = SecuritySeverity::Allow
    };
}

bool SecurityManager::is_connection_secure() const {
    // In a real implementation, this would check TLS/SSL status
    // For simulation, we'll just return the config setting
    return false; // Assume unsecure for simulation
}

bool SecurityManager::is_origin_allowed(const std::string& origin) const {
    if (impl_->config.allowed_origins.empty()) {
        return true; // No restriction
    }

    return impl_->config.allowed_origins.find(origin) != impl_->config.allowed_origins.end();
}

void SecurityManager::cleanup_expired_entries() {
    auto now = std::chrono::steady_clock::now();

    // Clean up every minute
    if (now - impl_->last_cleanup < std::chrono::minutes(1)) {
        return;
    }

    // Clean up expired IP blocks
    for (auto it = impl_->blocked_ips.begin(); it != impl_->blocked_ips.end();) {
        if (now > it->second) {
            it = impl_->blocked_ips.erase(it);
        } else {
            ++it;
        }
    }

    // Clean up expired rate limit entries
    for (auto it = impl_->rate_limits.begin(); it != impl_->rate_limits.end();) {
        if (now - it->second.window_start > std::chrono::seconds(impl_->config.rate_limit_window * 2)) {
            it = impl_->rate_limits.erase(it);
        } else {
            ++it;
        }
    }

    impl_->last_cleanup = now;
}

SecurityStats SecurityManager::get_stats() const {
    std::lock_guard<std::mutex> lock(impl_->security_mutex);

    SecurityStats stats;
    stats.blocked_ips = impl_->blocked_ips.size();
    stats.active_rate_limits = impl_->rate_limits.size();
    stats.config = impl_->config;
    stats.last_cleanup = impl_->last_cleanup;

    // Calculate total violations
    for (const auto& [_, info] : impl_->rate_limits) {
        stats.total_violations += info.violation_count;
    }

    return stats;
}

// Utility functions
SecurityConfig SecurityConfig::create_default() {
    return SecurityConfig{
        .require_authentication = false,
        .require_tls = false,
        .rate_limit_max_requests = 100,
        .rate_limit_window = 60, // 60 seconds
        .rate_limit_violations_before_block = 5,
        .ip_block_duration = 3600, // 1 hour
        .max_message_size = 1024 * 1024, // 1MB
        .min_token_length = 16,
        .max_token_length = 512,
        .token_validation = true
    };
}

SecurityConfig SecurityConfig::create_strict() {
    auto config = create_default();
    config.require_authentication = true;
    config.require_tls = true;
    config.rate_limit_max_requests = 10;
    config.rate_limit_window = 60;
    config.rate_limit_violations_before_block = 3;
    config.ip_block_duration = 7200; // 2 hours
    config.max_message_size = 64 * 1024; // 64KB
    config.min_token_length = 32;
    config.max_token_length = 256;

    // Add some common blocked patterns
    config.blocked_patterns = {
        "<script>",
        "javascript:",
        "eval(",
        "document.cookie",
        "localStorage",
        "sessionStorage"
    };

    return config;
}

SecurityConfig SecurityConfig::create_permissive() {
    auto config = create_default();
    config.require_authentication = false;
    config.require_tls = false;
    config.rate_limit_max_requests = 1000;
    config.rate_limit_window = 60;
    config.rate_limit_violations_before_block = 10;
    config.ip_block_duration = 300; // 5 minutes
    config.max_message_size = 10 * 1024 * 1024; // 10MB
    config.token_validation = false;

    return config;
}

std::string security_severity_to_string(SecuritySeverity severity) {
    switch (severity) {
        case SecuritySeverity::Allow:
            return "Allow";
        case SecuritySeverity::Throttle:
            return "Throttle";
        case SecuritySeverity::Authenticate:
            return "Authenticate";
        case SecuritySeverity::Encrypt:
            return "Encrypt";
        case SecuritySeverity::Block:
            return "Block";
        default:
            return "Unknown";
    }
}

std::string security_result_to_string(const SecurityResult& result) {
    std::stringstream ss;
    ss << "[" << security_severity_to_string(result.severity) << "] "
       << (result.allowed ? "ALLOWED" : "DENIED") << ": " << result.reason;
    return ss.str();
}

std::string security_stats_to_string(const SecurityStats& stats) {
    std::stringstream ss;
    ss << "Security Statistics:\n"
       << "  Blocked IPs: " << stats.blocked_ips << "\n"
       << "  Active Rate Limits: " << stats.active_rate_limits << "\n"
       << "  Total Violations: " << stats.total_violations << "\n"
       << "  Authentication Required: " << (stats.config.require_authentication ? "Yes" : "No") << "\n"
       << "  TLS Required: " << (stats.config.require_tls ? "Yes" : "No") << "\n"
       << "  Rate Limit: " << stats.config.rate_limit_max_requests
       << " requests per " << stats.config.rate_limit_window << " seconds";
    return ss.str();
}

// Random token generator
std::string generate_secure_token(size_t length) {
    static const char charset[] =
        "0123456789"
        "abcdefghijklmnopqrstuvwxyz"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "-_.";

    std::random_device rd;
    std::mt19937 rng(rd());
    std::uniform_int_distribution<size_t> dist(0, sizeof(charset) - 2);

    std::string token;
    token.reserve(length);
    for (size_t i = 0; i < length; ++i) {
        token += charset[dist(rng)];
    }

    return token;
}

} // namespace transport
} // namespace mcp